# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
rStar-Math 基础节点模块

这个模块定义了搜索树中节点的基础数据结构。BaseNode 是所有搜索算法节点的基类，
提供了树形结构的基本属性和操作方法。

主要功能：
1. 维护节点的状态信息（文本内容、额外信息等）
2. 管理树形结构的父子关系
3. 跟踪节点的深度和终止状态
4. 存储节点的奖励值和价值评估
5. 支持多模态数据存储和传递
"""
from __future__ import annotations
from typing import Optional, Dict, Any, List, Type
from pydantic import BaseModel, PrivateAttr, field_validator


class BaseNode(BaseModel):
    """
    搜索树基础节点类
    
    作为所有搜索算法节点的基类，提供树形结构的基本属性和操作。
    使用 Pydantic 进行数据验证和序列化支持。
    
    属性说明：
        state: 节点状态字典，存储文本内容和额外信息
        additional_state_keys: 额外的状态键列表，用于扩展状态信息
        parent: 父节点引用，根节点为 None
        children: 子节点列表
        depth: 节点在树中的深度，根节点为 0
        is_terminal: 是否为终止节点（叶子节点或达到终止条件）
        reward: 节点的奖励值，用于强化学习
        value: 节点的价值评估，用于搜索算法的决策
        tag: 节点的唯一标识符，用于调试和追踪
        consecutive_errors: 连续错误次数，用于错误处理
        multimodal_data: 多模态数据字典，存储图像等非文本信息
    """

    # 节点状态信息
    state: Dict[str, str] = {"text": "", "extra_info": ""}
    additional_state_keys: List[str] = []  # 额外的状态键
    
    # 树形结构属性
    parent: Optional[Any] = None  # 父节点
    children: List[Any] = []      # 子节点列表
    depth: int = 0
    
    # 节点状态标记
    is_terminal: bool = False     # 是否为终止节点
    
    # 价值评估相关
    reward: Optional[float] = None  # 奖励值
    value: Optional[float] = 0      # 价值评估

    # 节点标识和错误跟踪
    tag: str = "0"                # 节点唯一标识
    consecutive_errors: int = 0   # 连续错误次数
    
    # === 多模态数据支持 ===
    multimodal_data: Optional[Dict] = None  # 多模态数据字典，存储图像路径等信息

    def __init__(self, **kwargs) -> None:
        """
        初始化基础节点
        
        根据 additional_state_keys 初始化额外的状态字段，
        确保所有必要的状态键都存在于 state 字典中。
        
        参数:
            **kwargs: 传递给父类的初始化参数
        """
        super().__init__(**kwargs)

        # 为额外的状态键初始化空字符串值
        for key in self.additional_state_keys:
            self.state[key] = ""

    def has_children(self) -> bool:
        """
        检查节点是否有子节点
        
        返回:
            bool: 如果节点有子节点返回 True，否则返回 False
        """
        return self.children != []

    def is_root(self) -> bool:
        """
        检查节点是否为根节点
        
        返回:
            bool: 如果节点是根节点（没有父节点）返回 True，否则返回 False
        """
        return self.parent is None
